#!/usr/bin/env python
# -*- coding: UTF-8 -*-
"""
@Project ：machine_learning_warehouse 
@File ：神经网络分类问题处理框架.py
@IDE  ：PyCharm 
@Author ：李涵彬
@Date ：2024/11/29 16:50 
"""

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
import numpy as np


# 数据处理类
class DataProcessor:
	def __init__(self, X, y):
		self.X = X
		self.y = y

	def scale_data(self):
		scaler = StandardScaler()
		self.X = scaler.fit_transform(self.X)
		return self.X, self.y

	def train_test_split(self, test_size=0.2):
		X_train, X_test, y_train, y_test = train_test_split(self.X, self.y, test_size=test_size, random_state=42)
		return X_train, X_test, y_train, y_test

	def to_tensor(self):
		self.X = torch.tensor(self.X, dtype=torch.float32)
		self.y = torch.tensor(self.y, dtype=torch.long)
		return self.X, self.y


# 神经网络模型类
class NeuralNetwork(nn.Module):
	def __init__(self, input_size, hidden_size, num_classes):
		super(NeuralNetwork, self).__init__()
		self.fc1 = nn.Linear(input_size, hidden_size)
		self.relu = nn.ReLU()
		self.fc2 = nn.Linear(hidden_size, num_classes)

	def forward(self, x):
		out = self.fc1(x)
		out = self.relu(out)
		out = self.fc2(out)
		return out


# 训练和验证类
class Trainer:
	def __init__(self, model, criterion, optimizer, train_loader, valid_loader):
		self.model = model
		self.criterion = criterion
		self.optimizer = optimizer
		self.train_loader = train_loader
		self.valid_loader = valid_loader

	def train(self, num_epochs):
		for epoch in range(num_epochs):
			self.model.train()
			running_loss = 0.0
			for data, target in self.train_loader:
				self.optimizer.zero_grad()
				output = self.model(data)
				loss = self.criterion(output, target)
				loss.backward()
				self.optimizer.step()
				running_loss += loss.item()
			print(f'Epoch {epoch + 1}, Loss: {running_loss / len(self.train_loader)}')

			self.model.eval()
			with torch.no_grad():
				correct = 0
				total = 0
				for data, target in self.valid_loader:
					output = self.model(data)
					_, predicted = torch.max(output.data, 1)
					total += target.size(0)
					correct += (predicted == target).sum().item()
				accuracy = 100 * correct / total
				print(f'Epoch {epoch + 1}, Validation Accuracy: {accuracy}%')


# 主流程
def main():
	# 假设X和y是已经加载的数据和标签
	X, y = np.random.rand(100, 4), np.random.randint(0, 3, 100)  # 示例数据
	data_processor = DataProcessor(X, y)
	X, y = data_processor.scale_data()
	X_train, X_test, y_train, y_test = data_processor.train_test_split()
	X_train, y_train = data_processor.to_tensor()
	X_test, y_test = data_processor.to_tensor()

	train_dataset = TensorDataset(X_train, y_train)
	test_dataset = TensorDataset(X_test, y_test)
	train_loader = DataLoader(dataset=train_dataset, batch_size=10, shuffle=True)
	test_loader = DataLoader(dataset=test_dataset, batch_size=10, shuffle=False)

	model = NeuralNetwork(input_size=4, hidden_size=10, num_classes=3)
	criterion = nn.CrossEntropyLoss()
	optimizer = optim.Adam(model.parameters(), lr=0.001)

	trainer = Trainer(model, criterion, optimizer, train_loader, test_loader)
	trainer.train(num_epochs=10)


if __name__ == "__main__":
	main()
